from os import environ
from visdom import Visdom
import torch
import numpy as np
from tqdm import tqdm
from utils import *

wav_dir = './audio.worthsee.com/'
vis = Visdom()

def use_filter(wav_path : str):
    wav_name = wav_path.split("/")[-1]
    sample_rate, signal = wav_from_file(wav_path)
    mfcc = get_mfcc_from_array(signal, sample_rate)
    spectrum, freqs = get_fft_from_file(wav_path)
    filter_spectrum, filter_signal = filter_wav_from_spectrum(
        spectrum=spectrum,
        freqs=freqs,
        gain_factor=1.2,
        gain_range=(1000, 3000)
    )
    filter_signal = filter_signal.reshape([-1, 1])
    filter_signal = np.hstack((filter_signal, filter_signal))
    filter_mfcc = get_mfcc_from_array(filter_signal, sample_rate)

    env_name = "filter1000-3000"
    vis.heatmap(
        X=torch.tensor(mfcc), 
        win=wav_name + "mfcc without filter", 
        env=env_name,
        opts={
            "title" : "mfcc without filter\n{}".format(wav_name),
            "colormap" : "Viridis"  
        }
    )
    vis.heatmap(
        X=torch.tensor(filter_mfcc), 
        win=wav_name + "mfcc with filter", 
        env=env_name,
        opts={
            "title" : "mfcc with filter\n{}".format(wav_name),
            "colormap" : "Viridis"  
        }
    )

    indices = (freqs > 0)
    vis.line(
        X=torch.tensor(freqs[indices]),
        Y=torch.tensor(abs(spectrum[indices])),
        win=wav_name + "spectrum without filter",
        env=env_name,
        opts={
            "legend" : ["spectrum"], 
            "title" : "Frequency Domain(without filter)\n{}".format(wav_name),
            "color" : "dodgerblue"
        }
    )

    vis.line(
        X=torch.tensor(freqs[indices]),
        Y=torch.tensor(abs(filter_spectrum[indices])),
        win=wav_name + "spectrum with filter",
        env=env_name,
        opts={
            "legend" : ["spectrum"], 
            "title" : "Frequency Domain(with filter)\n{}".format(wav_name),
            "color" : "orangered"
        }
    )

    # display audio, remember to copy to avoid the warning
    vis.audio(
        tensor=torch.tensor(signal.copy()), 
        win=wav_path + ".audio without filter", 
        env=env_name, 
        opts={"title" : wav_path + " before filter", "sample_frequency" : sample_rate}
    )
    vis.audio(
        tensor=torch.tensor(filter_signal.copy()), 
        win=wav_path + ".audio with filter", 
        env=env_name, 
        opts={"title" : wav_path + " after filter", "sample_frequency" : sample_rate}
    )

def analyse_mfcc():
    for i in tqdm(range(1, 15)):
        wav_file = str(i) + ".wav"
        wav_path = wav_dir + wav_file

        sample_rate, signal = wav_from_file(wav_path)
        mfcc = get_mfcc_from_array(signal, sample_rate)

        vis.heatmap(
            X=torch.tensor(mfcc), 
            win=wav_file, 
            env="normal mfcc feature",
            opts={
                "title" : wav_file + "before",
                "colormap" : "Viridis"  
            }
        )

    for i in tqdm(range(121, 123)):
        wav_file = str(i) + "-空调异音.wav"
        wav_path = wav_dir + wav_file
        mfcc = get_mfcc_from_file(wav_path)
        vis.heatmap(
            X=torch.tensor(mfcc),
            win=wav_file,
            env="abnormal mfcc feature",
            opts={
                "title" : wav_file,
                "colormap" : "Viridis"
            }
        )


use_filter(wav_path=wav_dir+"1.wav")
use_filter(wav_path=wav_dir+"121-abnormal.wav")